using namespace std;

#include<iostream>
#include<fstream>
#include<sstream>
#include<cstdlib>
#include<gsl/gsl_rng.h>
#include<gsl/gsl_linalg.h>
#include<gsl/gsl_blas.h>
#include<gsl/gsl_randist.h>
#include"fast_svd.hpp"

fast_svd_t::fast_svd_t(int seed, ostream * os){
  this->os = os;
  rng = gsl_rng_alloc(gsl_rng_mt19937);
  gsl_rng_set(rng,seed);

}

fast_svd_t::~fast_svd_t(){
  gsl_rng_free(rng);
}

void fast_svd_t::debug_gsl_mat(const char * title, gsl_matrix * mat){
  *os<<"DEBUG MATRIX: "<<title<<endl;
  for(uint i=0;i<mat->size1;++i){
    for(uint j=0;j<mat->size2;++j){
      if(j) *os<<" ";
      *os<<gsl_matrix_get(mat,i,j);
    }
    *os<<endl;
  }
}

void fast_svd_t::debug_float_mat(const char * title, float * mat,int rows,int cols){
  *os<<"DEBUG MATRIX: "<<title<<endl;
  for(int i=0;i<rows;++i){
    for(int j=0;j<cols;++j){
      if(j) *os<<" ";
      *os<<mat[i*cols+j];
    }
    *os<<endl;
  }
}

void fast_svd_t::do_QR(float * floatA, int rows, int cols, float * floatQ, float * floatR){
  gsl_matrix * A = gsl_matrix_alloc(rows,cols);
  gsl_matrix * Q = gsl_matrix_alloc(rows,rows);
  gsl_matrix * R = gsl_matrix_alloc(rows,cols);
  *os<<"Copy A matrix into GSL object...\n";
  for(int i=0;i<rows;++i){
   for(int j=0;j<cols;++j){
     gsl_matrix_set(A,i,j,floatA[i*cols+j]);
   }
  }
  do_QR(A,Q,R);
  gsl_matrix_free(A);
  *os<<"Copy Q and R matrices into float array...\n";
  for(int i=0;i<rows;++i){
   for(int j=0;j<rows;++j){
     floatQ[i*rows+j] = gsl_matrix_get(Q,i,j);
   }
  }
  gsl_matrix_free(Q);
  for(int i=0;i<rows;++i){
   for(int j=0;j<cols;++j){
     floatR[i*cols+j] = gsl_matrix_get(R,i,j);
   }
  }
  gsl_matrix_free(R);
}

void fast_svd_t::do_QR(gsl_matrix * A,gsl_matrix * Q, gsl_matrix * R){
  int rows = A->size1;
  int cols = A->size2;
  int min = rows<=cols?rows:cols;
  gsl_vector * tau = gsl_vector_alloc(min);
  *os<<"QR decomposition begin...\n";
  os->flush();
  int rc1 = gsl_linalg_QR_decomp(A,tau);
  int rc2 = gsl_linalg_QR_unpack(A,tau,Q,R);
  if(rc1!=0 || rc2!=0){
    cerr<<"Error in QR decomposition.\n";
    exit(1);
  }
  gsl_vector_free(tau);
}

void fast_svd_t::do_SVD(float * floatA, int rows, int cols, float * floatU, float * floatS, float *floatV){
  gsl_matrix * A = gsl_matrix_alloc(rows,cols);
  *os<<"Copy A matrix into GSL object...\n";
  for(int i=0;i<rows;++i){
    for(int j=0;j<cols;++j){
      gsl_matrix_set(A,i,j,floatA[i*cols+j]);
    }
  }
  int u_rows,u_cols,v_rows,v_cols,s_cols;
  if(rows>=cols){
    u_rows = rows;
    u_cols = cols;
    s_cols = cols;
    v_rows = cols;
    v_cols = cols;
  }else{
    u_rows = rows;
    u_cols = rows;
    s_cols = rows;
    v_rows = cols;
    v_cols = rows;
  }
  gsl_matrix * U = gsl_matrix_alloc(u_rows,u_cols);
  gsl_vector * S =  gsl_vector_alloc(s_cols);
  gsl_matrix * V = gsl_matrix_alloc(v_rows,v_cols);
  do_SVD(A,U,S,V);
  gsl_matrix_free(A);
  *os<<"Copy U S and V matrices into float arrays...\n";
  for(int i=0;i<u_rows;++i){
    for(int j=0;j<u_cols;++j){
      floatU[i*u_cols+j] = gsl_matrix_get(U,i,j);
    }
  }
  gsl_matrix_free(U);
  for(int j=0;j<s_cols;++j){
    floatS[j] = gsl_vector_get(S,j);
  }
  gsl_vector_free(S);
  for(int i=0;i<v_rows;++i){
    for(int j=0;j<v_cols;++j){
      floatV[i*v_cols+j] = gsl_matrix_get(V,i,j);
    }
  }
  gsl_matrix_free(V);
}

void fast_svd_t::do_SVD(gsl_matrix * A, gsl_matrix * U,gsl_vector * S, gsl_matrix * V){
  int rc1;
  if(A->size1 >= A->size2){
    *os<<"Copying A matrix into U\n";
    for(uint i=0;i<A->size1 ; ++i){
      for(uint j=0;j<A->size2 ; ++j){
        gsl_matrix_set(U,i,j,gsl_matrix_get(A,i,j));
      }
    }
    gsl_vector * work = gsl_vector_alloc(S->size);
    rc1 = gsl_linalg_SV_decomp(U,V,S,work);
    gsl_vector_free(work);
  }else{
    gsl_matrix * U_copy = gsl_matrix_alloc(A->size2,A->size1);
    *os<<"Transposing A matrix into U\n";
    for(uint i=0;i<U_copy->size1;++i){
      for(uint j=0;j<U_copy->size2;++j){
        gsl_matrix_set(U_copy,i,j,gsl_matrix_get(A,j,i));
      }
    }
    gsl_matrix * V_copy = gsl_matrix_alloc(U_copy->size2,U_copy->size2);
    gsl_vector * work = gsl_vector_alloc(U_copy->size2);
    rc1 = gsl_linalg_SV_decomp(U_copy,V_copy,S,work);
    gsl_vector_free(work);
    *os<<"Copying V matrix into U\n";
    for(uint i=0;i<U->size1;++i){
      for(uint j=0;j<U->size2;++j){
        gsl_matrix_set(U,i,j,gsl_matrix_get(V_copy,i,j));
      }
    }
    gsl_matrix_free(V_copy);
    *os<<"Copying U matrix into V\n";
    for(uint i=0;i<V->size1;++i){
      for(uint j=0;j<V->size2;++j){
        gsl_matrix_set(V,i,j,gsl_matrix_get(U_copy,i,j));
      }
    }
    gsl_matrix_free(U_copy);
  }
  if(rc1!=0) {
    cerr<<"Error in SVD\n";
    exit(1);
  }

}


void fast_svd_t::do_randSVD(float * mat, int rows, int cols, float * floatU, float * floatS, float *floatV,int k,int p){
  // set up a GSL copy of the matrix
  gsl_matrix * A = gsl_matrix_alloc(rows,cols);
  *os<<"Copying float matrix into A\n";
  for(int i=0;i<rows;++i){
    for(int j=0;j<cols;++j){
      gsl_matrix_set(A,i,j,mat[i*cols+j]); 
    }
  }
  debug_gsl_mat("A",A);
  // set up a perturbation matrix 
  int perturb_cols = k+p;
  gsl_matrix * perturb = gsl_matrix_alloc(cols,perturb_cols);
  for(int i=0;i<cols;++i){
    for(int j=0;j<perturb_cols;++j){
      gsl_matrix_set(perturb,i,j,gsl_ran_gaussian(rng,1));
    }
  }
  debug_gsl_mat("perturb",perturb);
  gsl_matrix * y = gsl_matrix_alloc(rows,perturb_cols);
  gsl_blas_dgemm(CblasNoTrans,CblasNoTrans,1.0,A,perturb,0.0,y);
  gsl_matrix_free(perturb);
  debug_gsl_mat("y",y);
  gsl_matrix * Q = gsl_matrix_alloc(rows,rows);
  gsl_matrix * R = gsl_matrix_alloc(rows,perturb_cols);
  do_QR(y,Q,R);
  debug_gsl_mat("Q high rank",Q);
  gsl_matrix * Q_lowrank = gsl_matrix_alloc(rows,perturb_cols);
  for(int i=0;i<rows;++i){
    for(int j=0;j<perturb_cols;++j){
      double e = gsl_matrix_get(Q,i,j);
      gsl_matrix_set(Q_lowrank,i,j,e);
    }
  }
  gsl_matrix_free(y);
  gsl_matrix_free(Q);
  gsl_matrix_free(R);
  debug_gsl_mat("Q low rank",Q_lowrank);
  gsl_matrix * bt = gsl_matrix_alloc(cols,perturb_cols);
  gsl_blas_dgemm(CblasTrans,CblasNoTrans,1.0,A,Q_lowrank,0.0,bt);
  gsl_matrix_free(A);
  debug_gsl_mat("b",bt);
  int u_rows = cols;
  int u_cols = perturb_cols;
  int s_cols = perturb_cols;
  int v_rows = perturb_cols;
  int v_cols = perturb_cols;
  gsl_matrix * U = gsl_matrix_alloc(u_rows,u_cols);
  gsl_vector * S = gsl_vector_alloc(s_cols);
  gsl_matrix * V_small = gsl_matrix_alloc(v_rows,v_cols);
  do_SVD(bt,U,S,V_small);
  gsl_matrix_free(bt);
  debug_gsl_mat("U",U);
  debug_gsl_mat("V_small",V_small);
  gsl_matrix * V = gsl_matrix_alloc(rows,perturb_cols);
  gsl_blas_dgemm(CblasNoTrans,CblasNoTrans,1.0,Q_lowrank,V_small,0.0,V);
  gsl_matrix_free(Q_lowrank);
  gsl_matrix_free(V_small);
  // do a sanity check
  gsl_matrix * vs = gsl_matrix_alloc(V->size1,V->size2);
  for(uint i=0;i<V->size1;++i){
    for(uint j=0;j<V->size2;++j){
      double d = gsl_matrix_get(V,i,j) * gsl_vector_get(S,j);
      gsl_matrix_set(vs,i,j,d);
    }
  }
  *os<<"Running sanity check...\n";
  gsl_matrix * vsu = gsl_matrix_alloc(rows,cols);
  gsl_blas_dgemm(CblasNoTrans,CblasTrans,1.0,vs,U,0.0,vsu);
  debug_gsl_mat("USV:",vsu);
  gsl_matrix_free(vs);
  gsl_matrix_free(vsu);
  *os<<"Copying U and V into float matrices\n";
  for(uint i=0;i<V->size1;++i){
    for(uint j=0;j<V->size2;++j){
      floatU[i*V->size2+j] = gsl_matrix_get(V,i,j);
    }
  }
  for(int i=0;i<s_cols;++i) floatS[i] = gsl_vector_get(S,i);
  for(uint i=0;i<U->size1;++i){
    for(uint j=0;j<U->size2;++j){
      floatV[i*V->size2+j] = gsl_matrix_get(U,i,j);
    }
  }
  gsl_vector_free(S);
  gsl_matrix_free(U);
  gsl_matrix_free(V);
}


void test_QR(){
  int rows = 2;
  int cols = 3;
  int seed = 1;
  float testmat[] = {1,2,3,4,5,6};
  fast_svd_t fast_svd(seed,&cerr);
  float Q[rows*rows];
  float R[rows*cols];
  fast_svd.do_QR(testmat,rows,cols,Q,R);
  fast_svd.debug_float_mat("Q",Q,rows,rows);
  fast_svd.debug_float_mat("R",Q,rows,cols);
}


void test_SVD(){
  int k = 0;
  int p = 1;
  int kp = k+p;
  //int rows = 2;
  //int cols = 3;
  //float testmat[] = {1,2,3,4,5,6};
  int rows = 3;
  int cols = 2;
  float testmat[] = {1,4,2,5,3,6};
  int seed = 1;
  fast_svd_t fast_svd(seed,&cerr);
  int u_rows,u_cols,s_cols,v_rows,v_cols;
  bool do_rand = true;
  if(do_rand){
    u_rows = rows;
    u_cols = kp;
    s_cols = kp;
    v_rows = cols;
    v_cols = kp; 
  }else{
    if (rows>=cols){
      u_rows = rows;
      u_cols = cols;
      s_cols = cols;
      v_rows = cols;
      v_cols = cols;
    }else{
      u_rows = rows;
      u_cols = rows;
      s_cols = rows;
      v_rows = cols;
      v_cols = rows;
    }
  }
  float U[u_rows*u_cols];
  float S[u_cols];
  float V[u_cols*u_cols];
  if(do_rand){
    fast_svd.do_randSVD(testmat,rows,cols,U,S,V,0,1);
  }else{
    fast_svd.do_SVD(testmat,rows,cols,U,S,V);
  }
  fast_svd.debug_float_mat("U",U,u_rows,u_cols);
  fast_svd.debug_float_mat("S",S,s_cols,1);
  fast_svd.debug_float_mat("V",V,v_rows,v_cols);
}

#ifdef _UNIT_TEST
int main(int argc,char * argv[]){
  //test_QR();
  test_SVD();
  return 0;
}
#endif
